#  -*- coding: UTF-8 -*-
import matplotlib.pyplot as plt
import numpy as np
from pandas import Categorical
import operator
from collections import Counter


def createDataSet():
    dataSet = np.array([[1, 1, 'yes'], [1, 1, 'yes'], [1, 0, 'no'],
                        [0, 1, 'no'], [0, 1, 'no']])
    labels = Categorical(['no surfacing', 'flippers'])
    # change to discrete values
    return dataSet, labels


# 功能：计算 array-n*m数据集的column列的香农熵
def calcShannonEnt(dataSet, column=-1):
    # param:list-n*3,
    # return:float
    returnData = -1.0
    labels = Categorical(dataSet[:, column])  # 用-1能增加函数的可扩展性
    px = labels.describe().freqs
    entVector = -px * np.log2(px)
    returnData = entVector.sum()
    return returnData


# 功能：分割数据
def splitDataSet(dataSet, featureIdx, value):
    # param:array-n*m,int,element
    # return:subDataSet
    returnData = dataSet[dataSet[:, featureIdx] == value]
    return np.delete(returnData, featureIdx, 1)


# 功能：给定数据的特征，计算H（Y|X）
def calcHy_x(dataSet, column):
    # param:,array-n*m,int
    # return:float
    returnData = 0.0
    defineDomain = np.unique(dataSet[:, column])
    for x in defineDomain:
        filterData = splitDataSet(dataSet, column, x)
        Pxi = filterData.shape[0] / float(dataSet.shape[0])
        Hy_xi = calcShannonEnt(filterData)
        returnData += Pxi * Hy_xi

    return returnData


# 功能：给定数据集，选出最好的划分特征
def chooseBestFeature(dataSet):
    # param:data:array-n*m
    # return:column:int
    returnData = -1
    Hy = calcShannonEnt(dataSet)
    featureNum = dataSet.shape[1] - 1
    BestEntropy = -1.0
    for i in range(featureNum):
        Hy_x = calcHy_x(dataSet, i)
        InfoGain = Hy - Hy_x
        if InfoGain > BestEntropy:
            BestEntropy = InfoGain
            returnData = i
    return returnData


# 功能：
def createTree(dataSet, labels):
    # param:
    # return:
    returnData = {}
    classList = [x[-1] for x in dataSet]
    counterResult = Counter(classList)
    # 叶子节点
    if dataSet.ndim == 1:
        return max(counterResult, key=counterResult.get)
    if counterResult[classList[0]] == len(classList):
        return classList[0]
    # 分支节点
    bestFeatureIdx = chooseBestFeature(dataSet)
    defineDomain = np.unique(dataSet[:, bestFeatureIdx])
    nodeName = labels[bestFeatureIdx]
    returnData = {nodeName: {}}
    subLabels = labels.remove_categories(nodeName).dropna()
    for branchName in defineDomain:
        subData = splitDataSet(dataSet, bestFeatureIdx, branchName)
        returnData[nodeName][branchName] = createTree(subData, subLabels)
    return returnData


def plotNode(nodeTxt, centerPt, parentPt, nodeType):
    createPlot.ax1.annotate(nodeTxt,
                            xy=parentPt,
                            xycoords='axes fraction',
                            xytext=centerPt,
                            textcoords='axes fraction',
                            va="center",
                            ha="center",
                            bbox=nodeType,
                            arrowprops=arrow_args)


# 功能：
def createPlot():
    # param:
    # return:
    returnData = []
    fig = plt.figure(1, facecolor='white')
    fig.clf()  #清除坐标轴
    createPlot.ax1 = plt.subplot(111, frameon=False)
    plotNode('dn', (0.5, 0.1), (0.1, 0.5), decisionNode)
    plt.show()
    return returnData


# debug：calcShannonEnt,
dataSet, labels = createDataSet()
returnData = {}
# content:
decisionNode = {"boxstyle": "sawtooth", 'fc': "0.8"}
leafNode = {"boxstyle": "round4", 'fc': "0.8"}
arrow_args = {'arrowstyle': '<-'}
createPlot()
# return:
result = 0
print 'return', returnData, result

# -------------------